#include <flashai/commons/spdlog_utils.h>
#include <flashai/commons/statistics_timer.h>
#include <flashai/commons/string_utils.h>
#include <flashai/model_helper/model_helper_factory.h>
#include <ros/ros.h>

#include <chrono>
#include <thread>

#include "flashros/commons/utils.h"
#include "flashros/inference/pipeline_factory.h"

class FlashROSDetectorRunner {
 public:
  FlashROSDetectorRunner(int argc, char** argv) { Init(argc, argv); }

  void Init(int argc, char** argv) {
    // const params
    std::string inference_section = "inference";
    std::string data_root_ros_param_name = "data_root";
    std::string base_config_name = "/config_filename";

    // ros init
    flashros::commons::RosInit(argc, argv, "default_detection_inference");
    nh_ = new ros::NodeHandle;
    std::string node_name = ros::this_node::getName();

    // get data root
    std::string data_root =
        flashros::commons::GetRosParam(*nh_, data_root_ros_param_name);
    ROS_WARN_STREAM("Data root " << data_root);

    // Read config file
    std::string config_file_path =
        flashros::commons::GetRosParam(*nh_, node_name + base_config_name);
    flashai::utils::CheckFileExists(config_file_path);
    cfg_reader_ = new INIReader(config_file_path);
    ROS_WARN_STREAM("config file path " << config_file_path);

    // Init spdlog
    flashai::utils::InitSpdLogger(*cfg_reader_);
    SPDLOG_INFO("{} successfully init spdlog", node_name);

    // init flashai detector
    timer_.Tic("detector_init");
    std::vector<std::string> tokens;
    flashai::utils::SplitStr(
        cfg_reader_->GetString(inference_section, "model_sections", ""),
        tokens);
    for (auto& cur_section : tokens) {
      detectors_.push_back(
          flashai::model::ModelHelperFactory::GetInstance().GetDetectionHelper(
              *cfg_reader_, cur_section, data_root));
    }
    SPDLOG_INFO("Successfully init FlashAI Detector");
    timer_.Tac("detector_init");

    // init input and output
    timer_.Tic("input_output_pipeline_init");
    input_ = flashros::inference::InputPipelineFactory::GetInstance()
                 .GetInputPipeline(*cfg_reader_, *nh_, data_root);
    output_ = flashros::inference::OutputPipelineFactory::GetInstance()
                  .GetDetectionOutputInterface(*cfg_reader_, *nh_, data_root);
    timer_.Tac("input_output_pipeline_init");

    // configs for max inference speed
    int max_inference_fps =
        cfg_reader_->GetInteger(inference_section, "max_inference_fps", 10);
    min_inference_ms_iter_ = (int)(1000 / max_inference_fps);

    // 实现根据 class_id 过滤box
    tokens.clear();
    std::string keep_ids_strs =
        cfg_reader_->GetString(inference_section, "keep_ids", "");
    flashai::utils::SplitStr(keep_ids_strs, tokens, ":");
    assert(tokens.size() == detectors_.size());
    keep_ids_.resize(tokens.size());
    for (int i = 0; i < tokens.size(); i++) {
      std::vector<std::string> splits;
      flashai::utils::SplitStr(tokens[i], splits, ",");
      for (auto& id_str : splits) {
        keep_ids_[i].insert(std::stoi(id_str));
      }
      if (keep_ids_[i].size() > 0) {
        SPDLOG_INFO(
            "Successfully init {} keep ids for the {} detector, which is/are "
            "{}",
            keep_ids_[i].size(), i + 1, tokens[i]);
      }
    }

    // 实现根据 id 映射
    tokens.clear();
    std::string map_ids_str =
        cfg_reader_->GetString(inference_section, "map_ids", "");
    flashai::utils::SplitStr(map_ids_str, tokens, ":");
    assert(tokens.size() == detectors_.size());
    map_ids_.resize(tokens.size());
    for (int i = 0; i < tokens.size(); i++) {
      assert(tokens.size() % 2 == 0);
      std::vector<std::string> cur_tokens;
      flashai::utils::SplitStr(tokens[i], cur_tokens, ",");
      for (int j = 0; j < cur_tokens.size(); j += 2) {
        map_ids_[i][std::stoi(cur_tokens[j])] = std::stoi(cur_tokens[j + 1]);
      }
      if (map_ids_[i].size() > 0) {
        SPDLOG_INFO("Successfully init {} pairs of map ids, which is/are {}",
                    map_ids_[i].size(), tokens[i]);
      }
    }

    // 统计初始化时间
    timer_.Statistics();
    timer_.Clear();
  }

  void Run() {
    // prepare for inference
    int iter_cnt = 0;
    const int statistics_cnt =
        cfg_reader_->GetInteger("inference", "statistics_cnt", 100);

    while (nh_->ok()) {
      SPDLOG_DEBUG("Start Read frame from source");
      auto iter_start = std::chrono::system_clock::now();

      // Model Inference
      int duration = InferenceOnce();
      if (duration == -1) {
        SPDLOG_ERROR("Error reading frames from source...");
        break;
      }

      // print performance statistics every statistics_cnt steps
      iter_cnt++;
      if (iter_cnt == statistics_cnt) {
        iter_cnt = 0;
        timer_.Statistics();
        timer_.Clear();
      }
    }
  }

  int InferenceOnce() {
    timer_.Tic("iter");
    cv::Mat frame;

    timer_.Tic("read");
    bool ret = input_->Read(frame);
    if (!ret) return -1;
    timer_.Tac("read");

    timer_.Tic("detector-inference");
    std::vector<flashai::commons::Box> boxes;
    for (int i = 0; i < detectors_.size(); i++) {
      detectors_[i]->Inference(frame);
      std::vector<flashai::commons::Box> cur_boxes =
          detectors_[i]->GetPredictions();
      for (auto iter = cur_boxes.begin(); iter != cur_boxes.end();) {
        if (keep_ids_[i].size() > 0 &&
            keep_ids_[i].count((*iter).class_id) == 0) {
          // 如果不在 keep_ids 中，则删除该元素
          cur_boxes.erase(iter);
        } else {
          // 判断当前box的id是否需要转换
          if (map_ids_[i].count((*iter).class_id) > 0) {
            (*iter).class_id = map_ids_[i][(*iter).class_id];
          }
          iter++;
        }
      }
      SPDLOG_DEBUG("cur {} boxes, before total {}", cur_boxes.size(),
                   boxes.size());
      boxes.insert(boxes.end(), cur_boxes.begin(), cur_boxes.end());
      SPDLOG_DEBUG("cur {} boxes, after total {}", cur_boxes.size(),
                   boxes.size());
    }
    timer_.Tac("detector-inference");

    timer_.Tic("filter-map");
    timer_.Tac("filter-map");

    timer_.Tic("visualize");
    output_->Write(frame, boxes);
    timer_.Tac("visualize");

    int iter_duration = timer_.Tac("iter");
    std::string msg_format = "Successfully get {} boxes in {} ms";
    if (iter_duration >= min_inference_ms_iter_) {
      msg_format += ", won't sleep";
    } else {
      int sleep_ms = min_inference_ms_iter_ - iter_duration;
      msg_format += ", will sleep " + std::to_string(sleep_ms) + " ms";
      std::this_thread::sleep_for(std::chrono::milliseconds(sleep_ms));
    }
    SPDLOG_INFO(msg_format, boxes.size(), iter_duration);

    return iter_duration;
  }

  void Release() {
    for (auto& detector : detectors_) {
      detector->Release();
      delete detector;
    }
    detectors_.clear();
    input_->Release();
    delete input_;
    output_->Release();
    delete output_;
    delete cfg_reader_;
    timer_.Clear();
  }

 private:
  std::vector<
      flashai::model::BaseModelInterface<std::vector<flashai::commons::Box>>*>
      detectors_;
  flashai::utils::StatisticsTimer timer_;
  flashros::inference::InputInterface* input_;
  flashros::inference::DetectionOutputInterface* output_;
  ros::NodeHandle* nh_;
  INIReader* cfg_reader_;
  int max_wait_secs_ = 500000;
  int min_inference_ms_iter_;

  std::vector<std::set<int>> keep_ids_;
  std::vector<std::map<int, int>> map_ids_;
};

flashros::inference::InputInterface* GetInputPipeline(
    INIReader& reader, ros::NodeHandle& nh, std::string data_root = "") {
  std::string section_ = "input";
  auto input_pipeline_type =
      (flashros::inference::InputPipelineType)reader.GetInteger(section_,
                                                                "type", -1);
  switch (input_pipeline_type) {
    case flashros::inference::VIDEO_INPUT:
      return new flashros::inference::VideoCaptureInput(
          data_root + reader.GetString(section_, "video_path", ""), nh);
    case flashros::inference::ROS_INPUT:
      return new flashros::inference::RosInput(
          reader.GetString(section_, "camera_topic_name", ""),
          reader.GetInteger(section_, "max_wait_secs", 500), nh);
  }
  SPDLOG_ERROR("Input Pipeline Type {}", input_pipeline_type);
  return NULL;
}

int main(int argc, char** argv) {
  FlashROSDetectorRunner runner(argc, argv);
  runner.Run();
  runner.Release();
  return 0;
}